import tensorflow as tf

drop_out = 0.75


def conv2d(name, x, w_shape, strides=1):
    with tf.name_scope(name) as sc:
        w = tf.Variable(tf.truncated_normal(w_shape, stddev=1e-1))
        b = tf.Variable(tf.zeros(w_shape[3]))
    x = tf.nn.conv2d(x, w, [1, strides, strides, 1], padding='SAME')
    x = tf.nn.bias_add(x, b)
    x = tf.nn.relu(x, name=name)
    tf.summary.histogram(x.op.name, x)
    return x


def max_pool(x, name, k=2):
    return tf.nn.max_pool(x, ksize=[1, k, k, 1], strides=[1, k, k, 1], padding="SAME", name=name)


def alexnet(x):
    x = tf.reshape(x, shape=[-1, 28, 28, 1])
    conv1 = conv2d("conv1", x, w_shape=[11, 11, 1, 96])
    pool1 = max_pool(conv1, k=2, name="pool1")
    norm1 = tf.nn.lrn(pool1, depth_radius=4, bias=1.0, alpha=0.001 / 9.0, name="norm1")

    conv2 = conv2d("conv2", norm1, w_shape=[5, 5, 96, 256])
    pool2 = max_pool(conv2, k=2, name="pool2")
    norm2 = tf.nn.lrn(pool2, depth_radius=4, bias=1.0, alpha=0.001 / 9.0, name="norm2")

    conv3 = conv2d("conv3", norm2, w_shape=[3, 3, 256, 384])
    pool3 = max_pool(conv3, k=2, name="pool3")
    norm3 = tf.nn.lrn(pool3, depth_radius=4, bias=1.0, alpha=0.001 / 9.0, name="norm3")

    conv4 = conv2d("conv4", norm3, w_shape=[3, 3, 384, 384])
    conv5 = conv2d("conv5", conv4, w_shape=[3, 3, 384, 256])
    pool5 = max_pool(conv5, k=2, name="pool5")
    norm5 = tf.nn.lrn(pool5, depth_radius=4, bias=1.0, alpha=0.001 / 9.0, name="norm5")

    with tf.name_scope("fc1") as fc1_scope:
        wfc1 = tf.Variable(tf.truncated_normal([2 * 2 * 256, 4096], stddev=1e-1))
        bfc1 = tf.Variable(tf.zeros(4096))
        fc1 = tf.reshape(norm5, [-1, 2 * 2 * 256])
        fc1 = tf.nn.relu(tf.add(tf.matmul(fc1, wfc1), bfc1))
        fc1 = tf.nn.dropout(fc1, drop_out)

    with tf.name_scope("fc2") as fc2_scope:
        wfc2 = tf.Variable(tf.truncated_normal([4096, 4096], stddev=1e-1))
        bfc2 = tf.Variable(tf.zeros(4096))
        fc2 = tf.reshape(fc1, [-1, 4096])
        fc2 = tf.nn.relu(tf.add(tf.matmul(fc2, wfc2), bfc2))
        fc2 = tf.nn.dropout(fc2, drop_out)

    with tf.name_scope("out") as out_scope:
        wout = tf.Variable(tf.truncated_normal([4096, 10], stddev=1e-1))
        bout = tf.Variable(tf.zeros(10))
        out = tf.add(tf.matmul(fc2, wout), bout)
    return out
